
from torch.utils.data import Dataset
from torchvision import datasets, transforms

class DatasetSplit(Dataset):
    def __init__(self, dataset, idxs):
        self.dataset = dataset
        self.idxs = list(idxs)

    def __len__(self):
        return len(self.idxs)

    def __getitem__(self, item):
        image, label = self.dataset[self.idxs[item]]
        return image, label

def load_data(args):
    # # 加载MNIST数据集
    transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
    dataset_train = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
    dataset_test = datasets.MNIST(root='./data', train=False, download=True, transform=transform)
    dataset_val = datasets.MNIST(root='./data', train=False, download=True, transform=transform)
    
    return dataset_train, dataset_test, dataset_val
